import pyrealsense2 as rs
import numpy as np
import cv2
import os
from datetime import datetime

def initialize_cameras():
    """初始化所有可用的RealSense相机"""
    ctx = rs.context()
    devices = ctx.query_devices()  # 获取所有连接的设备
    
    if len(devices) == 0:
        raise Exception("未检测到任何RealSense相机")
    
    pipelines = []
    configs = []
    
    print(f"检测到 {len(devices)} 个RealSense相机，正在初始化...")
    
    for i, device in enumerate(devices):
        pipeline = rs.pipeline(ctx)
        config = rs.config()
        
        # 获取设备序列号（用于区分不同相机）
        serial = device.get_info(rs.camera_info.serial_number)
        config.enable_device(serial)
        
        # 启用彩色流和深度流（可根据相机型号调整分辨率和帧率）
        config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
        config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
        
        # 启动管道
        pipeline.start(config)
        pipelines.append((pipeline, serial))
        print(f"相机 {i+1} 初始化完成 (序列号: {serial})")
    
    return pipelines

def capture_and_save(pipelines, output_dir="captures"):
    """从所有相机捕获图像并保存"""
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")  # 用时间戳区分不同批次的照片
    
    for i, (pipeline, serial) in enumerate(pipelines):
        try:
            # 等待一帧数据
            frames = pipeline.wait_for_frames()
            color_frame = frames.get_color_frame()
            depth_frame = frames.get_depth_frame()
            
            if not color_frame or not depth_frame:
                print(f"相机 {i+1} (序列号: {serial}) 获取帧失败")
                continue
            
            # 转换为numpy数组
            color_image = np.asanyarray(color_frame.get_data())
            depth_image = np.asanyarray(depth_frame.get_data())
            
            # 为深度图应用彩色映射（便于可视化）
            depth_colormap = cv2.applyColorMap(
                cv2.convertScaleAbs(depth_image, alpha=0.03),
                cv2.COLORMAP_JET
            )
            
            # 保存图像
            color_path = os.path.join(output_dir, f"cam_{i+1}_serial_{serial}_color_{timestamp}.png")
            depth_path = os.path.join(output_dir, f"cam_{i+1}_serial_{serial}_depth_{timestamp}.png")
            
            cv2.imwrite(color_path, color_image)
            cv2.imwrite(depth_path, depth_colormap)
            
            print(f"相机 {i+1} 图像已保存:")
            print(f"  彩色图: {color_path}")
            print(f"  深度图: {depth_path}")
            
        except Exception as e:
            print(f"相机 {i+1} (序列号: {serial}) 捕获失败: {str(e)}")

def main():
    try:
        # 初始化所有相机
        pipelines = initialize_cameras()
        
        # 捕获并保存图像
        capture_and_save(pipelines)
        
    except Exception as e:
        print(f"发生错误: {str(e)}")
        
    finally:
        # 关闭所有相机管道
        if 'pipelines' in locals():
            for pipeline, _ in pipelines:
                pipeline.stop()
            print("所有相机已关闭")

if __name__ == "__main__":
    main()

